
import torch
from PIL import Image
from cameras import VideoDetection
from torchvision.transforms import v2 as T
import cv2

from models import get_model_instance_segmentation

def get_transform(train):
    transforms = []
    transforms.append(T.ToImage())
    transforms.append(T.ToDtype(torch.float32, scale=True))
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)

model = get_model_instance_segmentation(2)

model.load_state_dict(torch.load("./best.pth"))

model.eval()

class EvalImage:
    def __init__(self, path, trans = None) -> None:
        self.img = Image.open(path).convert("RGB")
        self.trans = trans

    def __call__(self):
        if self.trans is None:
            raise ValueError
        
        img= self.trans(self.img)

        return [img]
    

# img = cv2.imread("./data/o.jpg")



# res = get_transform(False)(img)
# print(res)

device = "cuda" if torch.cuda.is_available() else "cpu"

model.to(device)

videoDetection = VideoDetection(model, path="./data/1.flv", device=device)

videoDetection()